#!usr/bin/python3
#-*- coding: utf-8 -*-

import sys, os, time, asyncio, socket, random, threading, requests

from io import BytesIO
from contextlib import closing
from urllib import request
from Crypto.Cipher import AES

def random_ip():
    a=random.randint(1,255)
    b=random.randint(1,255)
    c=random.randint(1,255)
    d=random.randint(1,255)
    return(str(a)+'.'+str(b)+'.'+str(c)+'.'+str(d))
def random_Agent():
    agents = [
    "Mozilla/5.0 (Windows NT 6.3; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.153 Safari/537.36",
    "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:30.0) Gecko/20100101 Firefox/30.0",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.75.14 (KHTML, like Gecko) Version/7.0.3 Safari/537.75.14",
    "Mozilla/5.0 (compatible; MSIE 10.0; Windows NT 6.2; Win64; x64; Trident/6.0)",
    'Mozilla/5.0 (Windows; U; Windows NT 5.1; it; rv:1.8.1.11) Gecko/20071127 Firefox/2.0.0.11',
    'Opera/9.25 (Windows NT 5.1; U; en)',
    'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; .NET CLR 1.1.4322; .NET CLR 2.0.50727)',
    'Mozilla/5.0 (compatible; Konqueror/3.5; Linux) KHTML/3.5.5 (like Gecko) (Kubuntu)',
    'Mozilla/5.0 (X11; U; Linux i686; en-US; rv:1.8.0.12) Gecko/20070731 Ubuntu/dapper-security Firefox/1.5.0.12',
    'Lynx/2.8.5rel.1 libwww-FM/2.14 SSL-MM/1.4.1 GNUTLS/1.2.9',
    "Mozilla/5.0 (X11; Linux i686) AppleWebKit/535.7 (KHTML, like Gecko) Ubuntu/11.04 Chromium/16.0.912.77 Chrome/16.0.912.77 Safari/535.7",
    "Mozilla/5.0 (X11; Ubuntu; Linux i686; rv:10.0) Gecko/20100101 Firefox/10.0 ",
    'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:73.0) Gecko/20100101 Firefox/73.0',
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/70.0.3538.110 Safari/537.36'
    ]
    return random.choice(agents)

def add_header(text=True):
    base = [
        ("Accept-Language", "zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7"), 
        ("Proxy-Connection", "keep-alive"), 
        ('User-Agent', random_Agent()), 
        ('X-Forwarded-For', random_ip())
    ]
    if not text:
        base.append(('Accept-Encoding', 'gzip, deflate, br'))
    
    opener = request.build_opener()
    opener.addheaders = base
    request.install_opener(opener)

def resource_name(url):
    name = url[url.index('//') + 2:].split('/', 1)[1].replace('/', '_')
    return name if '?' not in name else name.split('?', 1)[0]

def downloadfile_(url, filename, timeout=60, text=False):
    def callbackfunc(blocknum, blocksize, totalsize):
        '''回调函数
        @blocknum: 已经下载的数据块
        @blocksize: 数据块的大小
        @totalsize: 远程文件的大小
        '''
        return
        downsize=blocknum * blocksize
        if downsize >= totalsize:
            downsize=totalsize
        percent = 100.0 * downsize / totalsize
        if percent > 100:
            percent = 100
        s ="%s: %6.2f%%"%(filename, percent)+"====>"+"%.2f"%(downsize/1024/1024)+"M/"+"%.2f"%(totalsize/1024/1024)+"M \r"
        sys.stdout.write(s)
        sys.stdout.flush()
        if percent == 100:
            print('')
    socket.setdefaulttimeout(timeout)
    try:
        add_header(text)
        request.urlretrieve(url, filename, callbackfunc)
    except socket.timeout:
        count = 1
        while count <= 5:
            try:
                add_header(text)
                request.urlretrieve(url, filename, callbackfunc)  
                break
                print('超时: %s' % url)
            except socket.timeout:
                count += 1
async def downloadfile(url, filename, params={}, timeout=30, text=False):
    try:
        result = None
#        print('url: %s, filename: %s' % (url, filename))
        with closing(requests.get(url, allow_redirects=True, timeout=(30,timeout), headers=add_header(text), params=params, stream=True)) as response:
            if response.status_code == 200:
                chunk_size = 1024
                content_size = int(response.headers.get('content-length', 0))
                if content_size <= 0:
                    result = response.content
                else:
                    buf = BytesIO()
                    for data in response.iter_content(chunk_size=chunk_size):
                        buf.write(data)
                    buf.seek(0)
                    result = buf.read()
            else:
                print('wget %s %s | [status: %s, reason: %s]\n%s' % (url, params, response.status_code, response.reason, response.content))
                exit(0)
        if result:
            with open(filename, 'ab') as file:
                    file.write(result)
        else:
            raise BaseException('result is None')
    except BaseException as e:
        print('Error: %s' % e)
        raise e

def downloadtask(urls, savepath, retry=5):
    tasks = []
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    os.makedirs(savepath, exist_ok = True)
    urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url)))]
    count = 0
    while(len(urls) > 0):
        print('task of %s: %s' % (threading.current_thread().getName(), len(urls)))
        tasks = [downloadfile(url, '%s/%s' % (savepath, resource_name(url))) for url in urls]
        try:
            loop.run_until_complete(asyncio.wait(tasks))
        except KeyboardInterrupt as e:
            print(asyncio.Task.all_tasks())
            for task in asyncio.Task.all_tasks():
                print(task.cancel())
            loop.stop()
            # loop.run_forever()
            break
        except BaseException as e:
            break
        urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url)))]
        count += 1
        if count >= retry: break
#        finally:
#            loop.close()
    loop.close()

def checktask_status(task_tss, savepath):
    while True:
        try:
            flag = True
            status = ''
            sum = 0
            for index, urls in enumerate(task_tss):
                urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url)))]
                status += '%6d\t' % len(urls)
                sum += len(urls)
                if len(urls) > 0: flag = False
            print('task remain: %s\t | sum(%s)' % (status, sum), end = '\r')
            if flag: break
        except KeyboardInterrupt as e:
            break;
    print('')
            
def getm3u8key(baseurl, keys, savepath):
    if len(keys) <= 0: return
    beg_m = keys.index('METHOD=') + 7
    end_m = keys.index(',')
    beg_u = keys.index('URI=') + 4
    method = keys[beg_m:end_m]
    if keys.index('"') > 0 or keys.index("'") > 0:
        key = keys[beg_u+1:-1]
    else:
        key = keys[beg_u:]
    url = '%s/%s' % (baseurl, key)
    with request.urlopen(url) as f:
        if f.status == 200:
            val = f.read()
            with open('%s/%s' % (savepath, key), 'wb') as dst:
                dst.write(val)
            return val
        else:
            print('Status:', f.status, f.reason)

def downloadm3u8(url, savepath):
    os.makedirs(savepath, exist_ok = True)
    all_content = ''
    print('download m3u8 file: %s' % url)
    downloadtask([url], savepath)
    filepath = '%s/%s' % (savepath, resource_name(url))
    if not os.path.exists(filepath):
        raise BaseException("m3u8 file not exists: %s" % filepath)
    with open(filepath, 'r') as m3u8:
        all_content = m3u8.read()
    if len(all_content) <= 0:
        raise BaseException("Failed to download m3u8 file: \n%s" % all_content)
    with open('%s/%s' % (savepath, resource_name(url)), 'a') as file:
        file.write('\n#EXTURL:%s' % url.rsplit('/', 1)[0])
    if '\r\n' in all_content:
        file_line = all_content.split("\r\n")
    else:
        file_line = all_content.split("\n")
    if file_line[0] != "#EXTM3U":
        raise BaseException("it is not m3u8 file: \n%s" % all_content)
    tss = []
    keys = ''
    url = url.rsplit('/', 1)[0]
    for line in file_line:
        if '#EXT-X-KEY' in line:
            keys = line
        if '#EXTURL:' in line:
            url = line[8:]
        if "#EXT" in line or len(line.strip()) == 0: continue
        tss.append(line)
    tss_out = []
    for ts in tss:
        if '.m3u8' in ts:
            tss_, key, url = downloadm3u8('%s/%s' % (url, ts), savepath)
            for ts_ in tss_:
                tss_out.append(ts_)
        else:
            tss_out.append(ts)
        
    key = getm3u8key(url, keys, savepath)

    return (tss_out, key, url)
#enddef

def loadm3u8file(savepath, filename):
    all_content = ''
    filepath = '%s/%s' % (savepath, filename)
    print('loading m3u8 file: %s' % (filepath))
    if not os.path.exists(filepath): raise BaseException('file not exists: %s' % filepath)
    with open(filepath, 'r') as file:
        all_content = file.read()
    if '\r\n' in all_content:
        file_line = all_content.split("\r\n")
    else:
        file_line = all_content.split("\n")
    if file_line[0] != "#EXTM3U":
        raise BaseException("it is not m3u8 file: \n%s" % all_content)
    tss = []
    keys = ''
    for line in file_line:
        if '#EXT-X-KEY' in line:
            keys = line
        if '#EXTURL:' in line:
            url = line[8:]
        if "#EXT" in line or len(line.strip()) == 0: continue
        tss.append(line)
    tss_out = []
    for ts in tss:
        if '.m3u8' in ts:
            tss_, key, url = downloadm3u8('%s/%s' % (url, ts), savepath)
            for ts_ in tss_:
                tss_out.append(ts_)
        else:
            tss_out.append(ts)
        
    key = getm3u8key(url, keys, savepath)
    return (tss_out, key, url)
    
def download_batch_ts(tss, savepath):
    size = len(tss)
    pagesize = size // 10 + 1
    pool = []
    for i in range(0, size, pagesize):
        t = threading.Thread(target=downloadtask, name='%s-%s' % (downloadtask.__name__, i // pagesize), args=(tss[i:i+pagesize], savepath))
        pool.append(t)
        t.start()
    t = threading.Thread(target=checktask_status, name='%s-%s' % (checktask_status.__name__, 1), args=([tss[i:i+pagesize] for i in range(0, size, pagesize)], savepath))
    t.start()
    t.join()
    # checktask_status(task_tss, savepath)
    # for t in pool:
    #    t.join()
def check_files(tss, path):
    ts_files = ['%s/%s/%s' % (os.getcwd(), path, resource_name(ts)) for ts in tss]
    file_not_exists = []
    exists = 0
    not_exists = 0
    for index, file_ in enumerate(ts_files):
        if not os.path.exists(file_):
            file_not_exists.append(tss[index])
            not_exists += 1
        else:
            exists += 1
        print('check_files in %s: [total: %-8d| exists: %-8d| not_exists: %-8d]' % (path, len(ts_files), exists, not_exists), end = '\r')
    print('')
    return file_not_exists

def mergefiles(tss, path, filename, key, rm_src = True, cover = False):
    ts_files = ['%s/%s' % (os.getcwd(), item) for item in os.listdir(path) if item.endswith('.ts')]
    if len(ts_files) != len(tss):
        raise BaseException('download file is not complete: %s (%s/%s)' % (path, len(ts_files), len(tss)))
    ts_files = ['%s/%s/%s' % (os.getcwd(), path, resource_name(ts)) for ts in tss]
    if len(ts_files) == 0: raise BaseException('num of ts_files is 0: %s' % ts_files)
    if len(filename) == 0: raise BaseException('filename is null: %s' % filename)
    file_not_exists = []
    for file_ in ts_files:
        if not os.path.exists(file_):
            file_not_exists.append(file_)
    if len(file_not_exists) > 0:
        raise BaseException('these file is not exists: %s' % file_not_exists)
    if os.path.exists(filename):
        if cover:
            os.remove(filename)
            print('cover dst file %s' % filename)
        else:
            filename = '%s_' % filename
    print('dst file save as: %s' % filename)
    with open(filename, 'ab') as dst:
        for index, file_ in enumerate(ts_files):
            with open(file_, 'rb') as src:
                if key and len(key) > 0:
                    cryptor = AES.new(key, AES.MODE_CBC, key)
                    dst.write(cryptor.decrypt(src.read()))
                else:
                    dst.write(src.read())
            print('%s: %s in %s | %3.2f%%' % ('mergefiles', index + 1, len(ts_files), ((index + 1) / len(ts_files) * 100.0)), end = '\r')
        print('')
    if rm_src:
        print('remove src files')
        for file_ in ts_files:
            os.remove(file_)
'''
转换加密文件
'''
def tansform_file(filename, key) :
    with open(filename, 'rb') as src:
        if key and len(key) > 0:
            cryptor = AES.new(key, AES.MODE_EAX, key)
            data = cryptor.decrypt(src.read())
            file_sub = filename.rsplit('.', 1)
            with open('%s_%s.%s' % (file_sub[0], key.decode('utf-8'), file_sub[1]), 'wb') as dst:
                dst.write(data)
'''
转换加密文件流
'''
def tansform_buf(buffer, key) :
    if key and len(key) > 0:
        cryptor = AES.new(key, AES.MODE_CBC, key)
        data = cryptor.decrypt(buffer.read())
        file_sub = filename.rsplit('.', 1)
        with open('%s_%s.%s' % (file_sub[0], key.decode('utf-8'), file_sub[1]), 'wb') as dst:
            dst.write(data)   
def load_m3u8key(baseurl, keys):
    if len(keys) <= 0: return
    beg_m = keys.index('METHOD=') + 7
    end_m = keys.index(',')
    beg_u = keys.index('URI=') + 4
    method = keys[beg_m:end_m]
    if keys.index('"') > 0 or keys.index("'") > 0:
        key = keys[beg_u+1:-1]
    else:
        key = keys[beg_u:]
    url = '%s/%s' % (baseurl, key)
    with request.urlopen(url) as rep:
        if rep.status == 200:
            buf = rep.read()
        else:
            buf = None
            print('Status:', f.status, f.reason)
    return buf

def process(text):
    sys.stdout.write('%s \r' % (text))
    sys.stdout.flush()
def main_file(path, m3u8):
    # 文件夹生成同名文件
    filename = path
    tss, key, url = loadm3u8file(path, m3u8)
    tss_not_exists = check_files(['%s/%s' % (url, ts) for ts in tss], path)
    download_batch_ts([ts for ts in tss_not_exists], path)
    mergefiles(['%s/%s' % (url, ts) for ts in tss], path, '%s.mp4' % filename, key)

def main_url(url, filename):
    # 生成同名文件夹
    path = filename
    tss, key, url = downloadm3u8(url, path)
    tss_not_exists = check_files(['%s/%s' % (url, ts) for ts in tss], path)
    download_batch_ts([ts for ts in tss_not_exists], path)
    mergefiles(['%s/%s' % (url, ts) for ts in tss], path, '%s.mp4' % filename, key)
if __name__ == '__main__':
    help_msg = '''\
        Usage:python3 %s -u url -p path -f filename
        -u: download url
        -p: dst save as <path>.mp4
        -f: load m3u8 file <path/filename>
        require pycryptodome, requests
        pip3 install pycryptodome requests -i https://pypi.tuna.tsinghua.edu.cn/simple
        '''    % sys.argv[0]
    if len(sys.argv) < 2:
        print(help_msg)
    else:
        start_t = time.time()
        # main_file
        if '-f' in sys.argv:
            filename = sys.argv[sys.argv.index('-f') + 1]
            if '-p' in sys.argv:
                path = sys.argv[sys.argv.index('-p') + 1]
            else:
                path, filename = filename.rsplit('/', 1)
            m3u8 = '%s/%s' % (path, filename)
            if not os.path.exists(m3u8):
                raise BaseException('file not exists: %s' % m3u8)
            main_file(path, filename)
            spend = int(time.time() - start_t)
            print('spend: %ds' % (spend))
            exit(0)
        # main_url
        if '-u' in sys.argv:
            url = sys.argv[sys.argv.index('-u') + 1]
            if '-p' not in sys.argv:
                path = resource_name(url)
                if '.' in path:
                    path = path.rsplit('.', 1)[0]
            else:
                path = sys.argv[sys.argv.index('-p') + 1]
            print('[DEBUG] path: %s' % path)
            main_url(url, path)
            spend = int(time.time() - start_t)
            print('spend: %ds' % (spend))
            exit(0)
        if '--test' in sys.argv:
            url = sys.argv[sys.argv.index('--test') + 1]
            name = resource_name(url)
            print(name)
            exit(0)
        if '--trans' in sys.argv:
            exit(0)
        print(help_msg)